# -*- coding: utf-8 -*-
"""

"""

import numpy as np
import copy


def softmax(x):
    probs = np.exp(x - np.max(x))
    probs /= np.sum(probs)
    return probs

class TreeNode(object):
    """A node in the MCTS tree.
Q是累积平均胜率，u则是对于Q的偏移,p先验概率

    """

    def __init__(self, parent, prior_p):
        self._parent = parent
        self._children = {}  # a map from action to TreeNode
        self._n_visits = 0
        self._Q = 0
        self._u = 0
        self._P = prior_p

    def expand(self, action_priors):
        """
            输入action_priors 是一个包括的所有合法动作的列表（list），表示在当前局面我可以在哪些地方落子。此函数为当前节点扩展了子节点
        """
        for action, prob in action_priors:
            if action not in self._children:
                self._children[action] = TreeNode(self, prob)

    def select(self, c_puct):
        """
        在子节中选择具有 （Q+u）最大的节点，c_puct是需要我们定义的值，在后面会说到。

        """
        return max(self._children.items(),
                   key=lambda act_node: act_node[1].get_value(c_puct))

    def update(self, leaf_value):


        self._n_visits += 1

        self._Q += 1.0*(leaf_value - self._Q) / self._n_visits

    def update_recursive(self, leaf_value):
        """
        该节点开始，自上而下地 更新 所有 的父节点
        """
        # If it is not root, this node's parent should be updated first.
        if self._parent:
            self._parent.update_recursive(-leaf_value)
        self.update(leaf_value)

    def get_value(self, c_puct):

        self._u = (c_puct * self._P *
                   np.sqrt(self._parent._n_visits) / (1 + self._n_visits))
        return self._Q + self._u

    def is_leaf(self):

        return self._children == {}

    def is_root(self):
        return self._parent is None


class MCTS(object):


    def __init__(self, policy_value_fn, c_puct=5, n_playout=10000):
        """
 policy_value_fn：当前采用的策略函数，输入是当前棋盘的状态，输出 (action, prob)元祖和score[-1,1]。
c_puct：控制探索和回报的比例，值越大表示越依赖之前的先验概率。
n_playout：MCTS的执行次数，值越大，消耗的时间越多，效果也越好
        """
        self._root = TreeNode(None, 1.0)
        self._policy = policy_value_fn
        self._c_puct = c_puct
        self._n_playout = n_playout
#一次迭代
    def _playout(self, state):
        """
        从根到叶运行单个播放，在叶上获取一个值，并通过其父级将其传播回来。状态已修改到位，因此必须提供副本
        """
        node = self._root
        while(1):
            if node.is_leaf():
                break

            action, node = node.select(self._c_puct)
            state.do_move(action)



        action_probs, leaf_value = self._policy(state)

        end, winner = state.game_end()
        if not end:
            node.expand(action_probs)
        else:

            if winner == -1:  # tie
                leaf_value = 0.0
            else:
                leaf_value = (
                    1.0 if winner == state.get_current_player() else -1.0
                )


        node.update_recursive(-leaf_value)

    def get_move_probs(self, state, temp=1e-3):
        """
        temp: temperature parameter in (0, 1] controls the level of exploration
        按顺序运行所有play，并返回可用的操作和它们相应的概率。
       状态：当前游戏状态
        temp：（0，1）中的温度参数控制探测级别
        """
        for n in range(self._n_playout):
            state_copy = copy.deepcopy(state)
            self._playout(state_copy)


        act_visits = [(act, node._n_visits)
                      for act, node in self._root._children.items()]
        acts, visits = zip(*act_visits)
        act_probs = softmax(1.0/temp * np.log(np.array(visits) + 1e-10))

        return acts, act_probs

    def update_with_move(self, last_move):
        """走到树上，保留我们已经知道的一切关于子树
        自我对弈时，每走一步之后更新MCTS的子树。
与玩家对弈时，每一个回合都要重置子树
        """
        if last_move in self._root._children:
            self._root = self._root._children[last_move]
            self._root._parent = None
        else:
            self._root = TreeNode(None, 1.0)

    def __str__(self):
        return "MCTS"


class MCTSPlayer(object):
    """AI player based on MCTS"""

    def __init__(self, policy_value_function,
                 c_puct=5, n_playout=2000, is_selfplay=0):
        self.mcts = MCTS(policy_value_function, c_puct, n_playout)
        self._is_selfplay = is_selfplay

    def set_player_ind(self, p):
        self.player = p

    def reset_player(self):
        self.mcts.update_with_move(-1)

    def get_action(self, board, temp=1e-3, return_prob=0):
        sensible_moves = board.availables

        move_probs = np.zeros(board.width*board.height)
        if len(sensible_moves) > 0:
            acts, probs = self.mcts.get_move_probs(board, temp)
            move_probs[list(acts)] = probs
            if self._is_selfplay:

                move = np.random.choice(
                    acts,
                    p=0.75*probs + 0.25*np.random.dirichlet(0.3*np.ones(len(probs)))
                )

                self.mcts.update_with_move(move)
            else:

                move = np.random.choice(acts, p=probs)

                self.mcts.update_with_move(-1)


            if return_prob:
                return move, move_probs
            else:
                return move
        else:
            print("WARNING: the board is full")

    def __str__(self):
        return "MCTS {}".format(self.player)
